import numpy as np
import scipy as sp
from scipy.optimize import minimize, check_grad
import pandas as pd
import sympy as smp

from IPython.display import display as ipydisplay
smp.init_printing()

ErrorTol = 0.05

class Family(object):

	def __init__(self,
		pi_nq = 0.2,
		pi_n = 0.6, pi_q = 0, pi_s = 1,
		alpha = 0.35, rho = -6, theta = 0.5, gamma = 1,
		y = 10,
		tu = 3):

		# set up parameter values
		self.a = np.array([alpha, rho, theta, gamma])
		self.pi = np.array([pi_n, pi_q, pi_s, pi_nq])

		self.y = y
		self.tu = tu

		# initial values of the full range of commodities
		self.xeset()

		# setup model
		self.setupModel()

		self.umprun = False
		self.dualOpt = False

	def xeset(self, xe = [1,1,1]):
		'''
		assign values to the full range of commodities
		do not it after __init__!!
		'''
		self.x = xe

	# Below are methods to set parameter values
	def xset(self, x = [1,1,1] ):
		'''
		assign values to the choice variables
		'''
		self.x = x
		# [n, q, s]
		self.bindParameters()
		self.updateNumerics()

	def yset(self, y = 10):
		'''
		assign a value to income
		'''
		self.y = y

	def tuset(self, tu = 3):
		'''
		assign a value to target utility level
		'''
		self.tu = tu

	def piset(self, pi = [1,1,1,0]):
		'''
		assign values to the price vector pi
		self.pi = np.array([pi_n, pi_q, pi_s, pi_nq])
		'''
		self.pi = np.array(pi)
		self.bindParameters()
		self.updateNumerics()

	def pifset(self, pif = [1,1,1]):
		'''
		assign values to fixed prices (excluding pi_nq)
		self.pi = np.array([pi_n, pi_q, pi_s, pi_nq])
		'''
		temp = list(pif)
		temp += list(self.pi[len(pif): ])
		self.pi = np.array(temp)

		self.bindParameters()
		self.updateNumerics()

	def aset(self, a = [0.5, 0, 0.5, 1] ):
		'''
		assign values to preference parameters
		self.a = np.array([alpha, rho, theta, gamma])
		'''
		self.a = np.array(a)
		self.bindParameters()
		self.updateNumerics()

	def prmtset(self, a = [0.5, 0, 0.5, 1],
		pi = [1,1,1,0],
		y = 10,
		tu = 3
		):
		'''
		assign values to parameters
		a =
		pi =
		y =
		tu =
		'''
		self.a = np.array(a)
		self.pi = np.array(pi)
		self.y = y
		self.tu = tu
		self.bindParameters()
		self.updateNumerics()

	def prmtset_pif(self, a = [0.5, 0, 0.5, 1],
		pif = [1,1,1],
		y = 10,
		tu = 3
		):
		'''
		assign values to parameters
		a =
		pif =
		y =
		tu =
		'''

		self.a = np.array(a)
		self.y = y
		self.tu = tu

		temp = list(pif)
		temp += list(self.pi[len(pif): ])
		self.pi = np.array(temp)
		#self.pifset(pif = pif) # automatic update model in self.pifset, but not setupmodel!

		self.bindParameters()
		self.updateNumerics()

	def setupModel(self):
		'''
		1. setup the symbolic representation of the model
		2. numerical evaluation of all functions
		'''
		# symbolic setup
		self.setupSymbols()

		# choice variables
		self.setupChoice()
		self.setupChoiceString() # for use in decomposition
		# D matrix
		self.setupSymbolicDerivatives()

		# which parameters to change
		self.bindParameters()
		# B matrix
		self.setupComparativeStatics()

		# numerical evaluations 
		self.updateNumerics()

	def setupSymbols(self):
		'''
		Symbolic model setup
		'''

		############################
		# the symbolic setup
		n, q, s = smp.symbols('n, q, s',
			positive = True, real = True)
		alpha, rho, theta, gamma  = \
			smp.symbols('alpha, rho, theta, gamma',
			real = True)
		pi_n, pi_q, pi_s = smp.symbols('pi_n, pi_q, pi_s',
			positive = True, real = True)
		pi_nq = smp.symbols('pi_nq', real = True)

		p_n, p_q, p_s = smp.symbols('p_n, p_q, p_s',
			positive = True, real = True)
		y, tu = smp.symbols('y u',
			positive = True, real = True)

		self.LAMDA, self.ETA = smp.symbols('lamda, eta',
			positive = True, real = True)

		self.A = smp.Matrix([alpha, rho, theta, gamma])
		self.PI = smp.Matrix([pi_n, pi_q, pi_s, pi_nq])
		self.PIF = smp.Matrix([pi_n, pi_q, pi_s]) # pif denotes the first three elements of pi

		self.XE = smp.Matrix([n, q, s]) # the full range of commodities

		self.Y = y
		self.TU = tu

		############################


		#############################
		# the utlity function

		rho_numeric = self.a[1]

		if rho_numeric == 0:
			u1 = n**alpha * (q**gamma)**(1-alpha)
		else:
			u1 = ( alpha*n**rho + (1-alpha)* (q**gamma)**rho
				 )**(1/rho)

		u = u1**theta * s**(1-theta)

		# store the symbol version of the utility function
		u = smp.expand_power_base(u)
		self.U = u

		################################
		# the expenditure function
		e = pi_n * n + pi_q * q + pi_s * s + pi_nq * n * q

		# store the symbolic version of the expenditure function
		self.E = e

	def setupChoice(self):
		'''
		Set up the choice variables
		'''
		# choice and parameters
		self.X = self.XE # X as choice variables

	def setupChoiceString(self):
		'''
		string representation of the choice variables
		for use in decomposition
		'''
		self.lenx = len(self.x)
		self.strx = [ str(self.X[ i ]) for i in range( self.lenx ) ]

	def setupSymbolicDerivatives(self):
		'''
		get the symbolic Jacobians, Hessians, and the D matrix

		D matrix: borded Hessian of choice variables
		D dx = B dbeta: dx = inv(D) B dbeta
		dbeta is the differentials of parameters
		'''
		# the symbolic utitily/expenditure Jacobians
		self.dU = smp.Matrix([self.U]).jacobian(self.X)
		self.dE = smp.Matrix([self.E]).jacobian(self.X)

		# the symbolic utility Hessian
		self.d2U = self.dU.jacobian(self.X)
		self.d2E = self.dE.jacobian(self.X)

		# the symbolic UMP D matrix

			# Dm: Hessian
		Dm = self.d2U - self.LAMDA*self.d2E
		
			# Dm + borders -> D: bordered Hessian
		tempDm = Dm.row_join(-self.dE.T)

		temprow = self.dE.row_join(smp.zeros(1))
			
		self.D = tempDm.col_join( - temprow )

		# the symbolic EMP D matrix

			# Dm: Hessian
		Dm = self.ETA * self.d2U - self.d2E

			# Dm + borders -> D: bordered Hessian
		tempDm = Dm.row_join(self.dU.T)

		temprow = self.dU.row_join(smp.zeros(1))

		self.Demp = tempDm.col_join( temprow )

	def bindParameters(self):
		'''
		assign numeric values of all parameters to math symbols
		'''
		# Parameter binding

		self.Aa = list(zip(self.A, self.a))
		self.PIpi = list(zip(self.PI, self.pi))

		self.Uprmt = self.Aa # parameters in the utility function
		self.Eprmt = self.PIpi + self.Aa # parameters in the expenditure function

		self.Dprmt = self.Aa + self.PIpi

		self.Bprmt = self.Dprmt

		# Define which parameters to change
		# BETA: price parameters that can change in the comparative static evaluations
		self.BETA  = self.PI # price parameters for comparative static evaluations

	def setupComparativeStatics(self):
		
		'''
		symbolic derivation of the B matrix
		'''
		# the symbolic UMP B matrix
		# 1. Given the FOCs on choice variables, take derivativs w.r.t. price parameters
		# 2. Move the derivatives of price parameters to the right-hand side
		
		# D dx = B dbeta: dx = inv(D) B dbeta
		# dbeta is the differentials of parameters

		# length of choice and price vectors
		templenx = len(self.X)
		templenb = len(self.BETA)

		Bm = self.LAMDA * self.dE.jacobian(self.BETA) - self.dU.jacobian(self.BETA)

		tempcol = smp.Matrix([0]* templenx )

		tempBm = Bm.row_join( tempcol)

		temprow = smp.Matrix([self.E]).jacobian(self.BETA)

		temprow = temprow.row_join( - smp.ones(1))

		s = [0]*templenb + [-1] # nullify the income effect
		temprow_y = smp.Matrix(s).T

		self.B = tempBm.col_join(temprow)

		self.B_y = tempBm.col_join(temprow_y) # no income effect

		# the symbolic EMP B matrix

		Bm = self.dE.jacobian(self.BETA) - self.ETA * self.dU.jacobian(self.BETA)

		tempcol = smp.Matrix([0]* templenx )

		tempBm = Bm.row_join( tempcol)

		temprow = smp.Matrix([ - self.U]).jacobian(self.BETA)

		temprow = temprow.row_join( smp.ones(1))

		s = [0]*templenb + [1] # nullify the real income effect
		temprow_u = smp.Matrix(s).T

		self.Bemp = tempBm.col_join(temprow)

		self.Bemp_u = tempBm.col_join(temprow_u)
		# no real income effect

	######################
	# evaluate functions
	######################
	def updateNumerics(self):
		self.updateFunctions()
		self.updateValues()
		self.updateConstraints()

	def updateFunctions(self):
		'''
		Update the numeric functions.
		Time consuming: 0.02 seconds
		'''

		### utility

		# the numeric utility function
		U_subs = self.U.subs(self.Uprmt)
		self.u_func = smp.lambdify(self.X, U_subs, 'numpy')

		# the numeric utility Jacobian
		dU_subs = self.dU.subs(self.Uprmt)
		self.du_func = smp.lambdify(self.X, dU_subs, 'numpy')

		# the numeric utility Hessian
		d2U_subs = self.d2U.subs(self.Uprmt)
		self.d2u_func = smp.lambdify(self.X, d2U_subs, 'numpy')

		### expenditure

		# the numeric expenditure function
		E_subs = self.E.subs(self.Eprmt)
		self.e_func = smp.lambdify(self.X, E_subs, 'numpy')

		# the numeric expenditure Jacobian
		dE_subs = self.dE.subs(self.Eprmt)
		self.de_func = smp.lambdify(self.X, dE_subs, 'numpy')

		# the numeric expenditure Hessian
		d2E_subs = self.d2E.subs(self.Eprmt)
		self.d2e_func = smp.lambdify(self.X, d2E_subs, 'numpy')

	def evalUtility(self, x, sign = 1.0):
		'''
		Evaluate the utility function
		'''
		u = self.u_func(*x)

		return sign * u

	def evalUjac(self, x, sign = 1.0):
		'''
		Evaluate the Jacobian of the utility function
		'''
		ujac = self.du_func(*x)

		return sign * np.asarray(ujac).ravel()

	def evalUhes(self, x, sign = 1.0):
		'''
		Evaluate the Hessian of the utility function
		'''
		uhes = self.d2u_func(*x)

		return sign * np.asarray(uhes)

	def evalExpenditure(self, x, sign = 1.0):
		'''
		Evaluate the expenditure function.
		'''
		e = self.e_func(*x)

		return sign * e

	def evalEjac(self, x, sign = 1.0):
		'''
		Evaluate the Jacobian of the expenditure function.
		'''
		ejac = self.de_func(*x)

		return sign * np.asarray(ejac).ravel()

	def evalEjac_m(self, x):
		'''
		Evaluate the negative Jacobian of the expenditure function.
		'''
		ejac = self.de_func(*x)

		return -1.0 * np.asarray(ejac).ravel()

	def evalEhes(self, x, sign = 1.0):
		'''
		Evaluate the Hessian of the expenditure function
		'''
		ehes = self.d2e_func(*x)

		return sign * np.asarray(ehes)

	def evalDB(self):
		'''
		Evaluate the D and B matrix of the UMP
		'''
		#print("Evaluate D B matrix ... ")
		self.Xx = list(zip(self.X, self.x))
		self.LAMDAlamda = [(self.LAMDA, self.lamda)]

		# D matrix
		s = self.Dprmt+ self.LAMDAlamda + self.Xx
		self.d = np.array( self.D.subs(s).evalf() ).astype(float)

		self.d_inv = np.linalg.inv(self.d)

		self.yeffect = -self.d_inv[:-1, -1:]

		self.slutsky = self.d_inv[ :-1, :-1]*self.lamda

		# B matrix

		s = self.Bprmt+ self.LAMDAlamda + self.Xx
		self.b = np.array( self.B.subs(s).evalf() ).astype(float)

		self.b_y = np.array( self.B_y.subs(s).evalf() ).astype(float) # no income effect

	def evalComparativeStatics(self):

		self.evalDB()

		self.cs = np.dot( self.d_inv, self.b )

		self.cs_y = np.dot( self.d_inv, self.b_y ) # no income effect
		# set the income effects to zero

	def evalDBemp(self):
		'''
		Evaluate the D and B matrix of the EMP
		'''

		self.Xx = list(zip(self.X, self.x))
		self.ETAeta = [(self.ETA, self.eta)]

		# D matrix
		s = self.Dprmt + self.ETAeta + self.Xx
		self.demp = np.array( self.Demp.subs(s).evalf() ).astype(float)

		self.demp_inv = np.linalg.inv(self.demp)

		self.ueffect = self.demp_inv[:-1, -1:]

		self.slutsky_emp = self.demp_inv[:-1,:-1]

		# B matrix
		s = self.Bprmt+ self.ETAeta + self.Xx
		self.bemp = np.array( self.Bemp.subs(s).evalf() ).astype(float)

		self.bemp_u = np.array( self.Bemp_u.subs(s).evalf() ).astype(float) # no income effect

	def evalComparativeStaticsEMP(self):

		self.evalDBemp()

		self.csemp = np.dot( self.demp_inv, self.bemp )

		self.csemp_u = np.dot( self.demp_inv, self.bemp_u) # no income effect

	def updateValues(self):
		'''
		Upate the utility level, expenditure level, and the actual prices.
		'''
		self.updateUtility()
		self.updateExpenditure()
		self.updateActualPrice()
		self.updateMarginalUtility()

	def updateUtility(self):
		'''
		Update the utility level.
		'''
		self.u = self.evalUtility(self.x)

	def updateExpenditure(self):
		'''
		Update the expenditure level.
		'''
		self.e = self.evalExpenditure(self.x)

	def updateActualPrice(self):
		'''
		Update the actual prices (the marginal expenditure).
		'''
		self.p = self.evalEjac(self.x)

	def updateMarginalUtility(self):
		'''
		Update the marginal utilities.
		'''
		self.mu = self.evalUjac(self.x)

	######################
	# optimizations
	######################

	def excessWealth(self, x):
		return self.y - self.evalExpenditure(x)

	def excessUtility(self, x):
		return self.evalUtility(x) - self.tu

	def updateConstraints(self):
		'''
		Update the constraints for the UMP and the EMP.
		'''
		self.cons_ump = (
			{
			'type':'eq',
			'fun': self.excessWealth,
			'jac': self.evalEjac_m
			},
			{
			'type':'ineq',
			'fun': lambda x: np.array(x),
			'jac': lambda x: np.identity(len(x))
			}
			)

		self.cons_emp = (
			{'type':'eq',
			'fun': self.excessUtility,
			'jac': self.evalUjac
			},
			{
			'type':'ineq',
			'fun': lambda x: np.array(x),
			'jac': lambda x: np.identity(len(x))
			}
			)

	def ump_simple(self, disp = False):
		'''
		The Utility Maximization Problem, simplified precedure
		'''
		if disp:
			print('-'*20)
			print('UMP Setup')
			print('Wealth: ' + str(self.y))
			print('Initial Commodity bunlde: ' + str(self.x))

		res = minimize(
			self.evalUtility, self.x, args = (-1,),
			jac = self.evalUjac,
			constraints = self.cons_ump,
			method = 'SLSQP',
			options = {'disp':disp}
			)

		if res.success: 
			self.x = list(res.x)
			self.u = -res.fun
			if disp:
				print('UMP Results')
				print('Optimal Choice: {}'.format(res.x))
				print('Optimal Utility Level: {0}'.format(-res.fun) )
				print('-'*20)
		else: 
			print('UMP Failed! u = -9, x = [0,0]')
			print('res.x:', res.x)
			print('res.success:', res.success)
			print('-'*20)
			self.x = [0,0]
			self.u = -9

	def ump(self, disp = False):
		'''
		The Utility Maximization Problem
		'''
		if disp:
			print('-'*20)
			print('UMP Setup')
			print('Wealth: ' + str(self.y))
			print('Initial Commodity bunlde: ' + str(self.x))

		res = minimize(
			self.evalUtility, self.x, args = (-1,),
			jac = self.evalUjac,
			constraints = self.cons_ump,
			method = 'SLSQP',
			options = {'disp':disp}
			)
		if res.success: 
			self.x = list(res.x)
			self.u = -res.fun
			if disp:
				print('UMP Results')
				print('Optimal Choice: {}'.format(res.x))
				print('Optimal Utility Level: {0}'.format(-res.fun) )
				print('-'*20)
		else: 
			print('UMP Failed! u = -9, x = [0,0]')
			print('res.x:', res.x)
			print('res.success:', res.success)
			print('-'*20)
			self.x = [0,0]
			self.u = -9

		self.umpres = res

		if disp:
			print('Update values ... ')
		self.updateValues()

		# get the shadow price of income (lamda)
		self.lamda = np.mean( self.mu / self.p )

		# get the slutsky matrix, and the income effects
		if disp:
			print('Get the slutsky matrix, and the income effects ... ')
		self.evalComparativeStatics() # very time consuming

		self.umprun = True


	def emp(self, disp = False): # expenditure minimization
		'''
		The Expenditure Minimization Problem
		'''
		if disp:
			print('-'*20)
			print('EMP Information')
			print('Target Utility: ' + str(self.tu))
			print('Initial Commodity bunlde: ' + str(self.x))

		res = minimize(
			self.evalExpenditure, self.x, args = (1,),
			jac = self.evalEjac,
			constraints = self.cons_emp,
			method = 'SLSQP',
			options = {'disp':disp}
			)
		assert res.success, 'EMP Failed!'

		self.x = list(res.x)
		#print('EMP sucess? '+str(res.success))
		self.empres = res
		self.updateValues()
		self.eta = np.mean( self.p / self.mu )
		self.lamda = 1.0/self.eta

		self.evalComparativeStaticsEMP()

		if disp == True:
			print('Expenditure Minimization Problem!')
			print('Optimal Choice: {}'.format(res.x))
			print('Optimal Expenditure Level: {0}'.format(res.fun) )
			print('-'*20)

	def dualOptimize(self, umpfirst = True, disp = False):

		if umpfirst:
			self.ump(disp = disp)
			self.tuset(self.u)
			self.emp(disp = disp)
		else:
			self.emp(disp = disp)
			self.yset(self.e)
			self.ump(disp = disp)

		self.dualOpt = True


		self.testDuality(disp = disp)


	def testDuality(self, disp = False):

		error = np.mean( np.abs( self.slutsky - self.slutsky_emp ) )
		base = np.sum( np.abs( self.slutsky ) )


		if disp:
			print('-'*50)
			print('Duality Test in Dual Optimization (UMP+EMP):')

			if error < ErrorTol * base:
				print('Test Succeeded! Equivalence of Slutsky Confirmed')
				print('Error: ' + str(error))
			else:
				print('Duality Failed! See the two matrix:')
				print(self.slutsky)
				print(self.slutsky_emp)

			print('-'*50)

	########################
	# self demonstration
	########################

	def showSymbols(self):

		# show the symbolic setup of the model

		print('Choice Variables')
		ipydisplay(self.X)

		# utility
		print('Utility Function:')
		ipydisplay(self.U)

		print('Utility Jacobian')
		ipydisplay(self.dU)

		print('Utility Hessian')
		ipydisplay(self.d2U)

		print('Utility Parameters')
		ipydisplay(self.Uprmt)

		# expenditure
		print('Expenditure Function:')
		ipydisplay(self.E)

		print('Expenditure Jacobian')
		ipydisplay(self.dE)

		print('Expenditure Hessian')
		ipydisplay(self.d2E)

		print('Expenditure Parameters')
		ipydisplay(self.Eprmt)

		# D matrix
		print('Matrix D')
		ipydisplay(self.D)

	def show(self, symbol = False):

		print('-'*20)
		print('Current Family Information:')

		if symbol: self.showSymbols()

		s = 'Preference Parameters:\n' + '{:0.3f}  '*len(self.a)
		print(s.format(*self.a))

		s = 'Fixed Prices:\n' + '{:0.3f}  '*len(self.pi)
		print(s.format(*self.pi) )

		s = 'Commodity Bundle:\n' + '{:0.3f}  '*len(self.x)
		print(s.format(*self.x) )

		s = 'Marginal Utilities:\n' + '{:0.3f}  '*len(self.mu)
		print(s.format(*self.mu) )

		s = 'Actual Prices:\n' + '{:0.3f}  '*len(self.p)
		print(s.format(*self.p) )

		if self.umprun:
			print('Price Effects (Slutsky Matrix):')
			print(self.slutsky)
			print('Income Effects:')
			print(self.yeffect)

		print('Wealth: {:.3f}'.format(self.y) )
		print('Target Utility: {:.3f}'.format(self.tu) )
		print('Expenditure: {:.3f}'.format(self.e) )
		print('Utility: {:.3f}'.format(self.u) )
		print('-'*20)



class FamilyRation(Family):

	def setupChoice(self):
		'''
		Set up the choice variables
		'''
		# choice
		self.X = self.XE[1:, :] # X as choice variables
		self.N = self.XE[:1, :] # N as rationed choice variable

	def xeset(self, xe = [1,1,1]):
		'''
		do not use this method after __init__!!!
		'''
		# choice
		self.x = xe[1:]
		# parameter
		self.n = xe[:1]

	def xset(self, x = [1,1] ):
		self.x = x
		# [q, s]
		self.bindParameters()
		self.updateNumerics()

	def nset(self, n = [1]):
		self.n = n # n as list object

		# re-evaluate all numerical functions after setting a parameter
		self.bindParameters()
		self.updateNumerics()

	def getU2(self):

		FU2 = Family(pi_nq = self.pi[-1])
		FU2.prmtset_pif(
			a = self.a,
			pif = self.pi[:-1],
			y = self.y,
			tu = self.tu
			)

		return FU2

	def getR1(self):

		assert self.pi[-1] > 0, 'Can not get R1, pi_interact = 0'

		FR1 = FamilyRation(pi_nq = 0)
		FR1.prmtset_pif(
			a = self.a,
			pif = self.pi[:-1],
			y = self.y,
			tu = self.tu
			)
		FR1.nset(self.n)

		return FR1

	def getU1(self):

		FU1 = Family(pi_nq = 0)
		FU1.prmtset_pif(
			a = self.a,
			pif = self.pi[:-1],
			y = self.y,
			tu = self.tu
			)

		return FU1

	def dualOptimizeOffer(self, umpfirst = True, disp = False):
		self.dualOptimize(umpfirst = umpfirst, disp = disp)
		self.getShadowPriceN()
		self.getShadowPriceR1()
		self.getPiR1()
		self.getCS_N()

	def bindParameters(self):
		# parameters

		self.Aa = list(zip(self.A, self.a))
		self.PIpi = list(zip(self.PI, self.pi))

		self.Nn = list(zip(self.N, self.n))

		self.Uprmt = self.Aa + self.Nn # parameters in the utility function
		self.Eprmt = self.PIpi + self.Nn # parameters in the expenditure function

		self.Dprmt = self.Aa + self.PIpi + self.Nn
		self.Bprmt = self.Dprmt

		# BETA: the exogenous changes to consider
		temp = self.PI.col_join(self.N)

		self.BETA = temp

	def getShadowPriceN(self):
		'''
		From Restricted Model to Unrestricted Model
		Key attributes returned:

		spe (shadow price vector, entire): full range of shadow price (sp has less dimension than spe)

		spi (shadow fixed price): full range of shadow fixed price
		
		sy (short for "shadow income" in the 2015 version): income + compensating variation
		'''
		# get the shadow price of n
		self.dUdN = smp.Matrix( [self.U] ).jacobian(self.N)
		s = self.Uprmt + self.Xx
		self.dudn = np.array( self.dUdN.subs(s).evalf() ).astype(float)

		sp_n =self.dudn/self.lamda
		#print('sp_n')
		#print(sp_n)
		# shadow price, full price vector
		self.spe  = np.hstack( (sp_n.ravel(), self.p) )

		# get the actual price of n
		self.dEdN = smp.Matrix( [self.E] ).jacobian(self.N)
		s = self.Eprmt + self.Xx
		self.dedn = np.array( self.dEdN.subs(s).evalf() ).astype(float)

		p_n = self.dedn

		# actual price, full price vector
		self.pe = np.hstack( (p_n.ravel(), self.p) )
		#print('pe')
		#print(self.pe)
		# induced price change
		self.spe_pe = self.spe - self.pe # induced price change
		#print('spe_pe')
		#print(self.spe_pe)
		# induced fixed price change
		fixedlen = len(self.pe)

		pi_fixed = self.pi[ :fixedlen ]
		pi_interact = self.pi[ fixedlen: ]

		spi_f = pi_fixed + self.spe_pe # dot not need e to indicate full range

		self.spi = np.hstack( (spi_f, pi_interact) )

		# induced income change (compensating variation)
		ychange = np.dot( self.spe_pe, self.n + self.x )
		# caution: do not use self.xe after __init__!!!!

		self.sy = self.y + ychange # shadow income! haha ~

	def getShadowPriceR1(self):
		'''
		obsolete!!
		From q2 model to q1 model
		key attribute returned:
		spi_q1: shadow pi price in the induced q1 model
		'''

		# assert self.pi[-1] > 0, 'Why do this in the Q1 model?'

		pi_ration = self.pi[:1]
		fixedlen = len(self.x + self.n)
		pi_fixed = self.pi[ :fixedlen ]


		spif_q1 = np.hstack( (pi_ration, self.p) )

		self.spi_q1 = np.hstack( (spif_q1, [0]) )

	def getPiR1(self):
		'''
		From q2 model to q1 model
		key attribute returned:
		pi_q1: pi price in the induced q1 model
		'''

		#assert self.pi[-1] > 0, 'Why do this in the Q1 model?'

		pi_ration = self.pi[:1] # pi_n

		pif_r1 = np.hstack( (pi_ration, self.p) ) # pi_q = p_q, pi_s = p_s 

		self.pi_r1 = np.hstack( (pif_r1, [0]) )

	def getCS_N(self):
		self.dxdn = self.cs[: -1, len(self.PI): -1 ]
		self.dxdn_emp = self.csemp[: -1, len(self.PI): -1 ]

	def evalDiscreteN(self, disp = False):
		self.n_array = np.arange(1,7)
		len_n = len(self.n_array)
		self.u_array = np.zeros(len_n)
		self.q_array = np.zeros(len_n)
		
		if disp: print("Compute optimal fertility under discrete choices, n in 1 to 6 ")

		for i in range(0,len(self.n_array) ):
			
			if disp: print("Evaluate n = " + str(self.n_array[i]) )

			# set initial values of q0 and s0
			#q0 = self.a[2] * net_y / ( self.pi[3] * self.n[0] )
			#s0 = (1 - self.a[2]) * net_y / self.pi[2] 
			#self.xset([q0, s0])
			self.x = [0.1, 0.1] # an easy solution

			self.n = [ self.n_array[i] ]

				# update the derivatives for given values of q, s, n
			self.bindParameters()
			self.updateNumerics()

			# pi_n: pi[0]; pi_nq: pi[3]; pi_s: pi[2]
			# theta: a[2]

			# income net of expenses on n_bar
			net_y = self.y - self.pi[0] * self.n[0]
			if disp: print("Net y excluding expenses on n: " + str(net_y))
			if net_y < 0.10 * self.y:
				if disp: print("Net y < 0.10*y, pass")
				self.u_array[i] = -9 # negative u to denote negative net y
				continue

			# run ump
			self.ump_simple(disp = disp)
			#self.ump(disp = disp)

			self.u_array[i] = self.u
			self.q_array[i] = self.x[0]
		
		i_max = np.argmax(self.u_array)

		self.n_opt = [ self.n_array[i_max] ]
		self.q_opt = self.q_array[i_max]

		#self.nset( self.n_opt )

		return self.n_opt
	

	def evalQQ_DiscreteN(self):
		'''
		evaluate QQ effects in semi-elasticity form 
		after the execution of evalDiscreteN()
		based on the stored n_array and q_array
		'''

		# qq effect in level terms, n margins: 1->2, 2->3, 3->4, 4->5, 5->6
		qq_level = self.q_array[1:] - self.q_array[:-1]

		q_base = (self.q_array[1:] + self.q_array[:-1]) / 2

		self.qq_e_DiscreteN = qq_level/q_base

		return self.qq_e_DiscreteN

	def evalQQ_Slutsky(self, n_bar = 2):
		'''
		evaluate QQ effect at n_bar using the Slutsky matrix
		return QQ effect in semi-elasticity form
		'''

		self.nset([n_bar])

		self.ump()
		self.evalComparativeStatics()

		self.qq_e_Slutsky = self.cs[: -1, len(self.PI): -1 ][0][0] / self.x[0]

		# return qq effect in semi-elasticity form
		return self.qq_e_Slutsky

	def evalQQ_DiscreteN(self):
		'''
		evaluate QQ effects in semi-elasticity form 
		after the execution of evalDiscreteN()
		based on the stored n_array and q_array
		'''

		# qq effect in level terms, n margins: 2->3, 3->4, 4->5
		qq_level = self.q_array[1:] - self.q_array[:-1]

		q_base = (self.q_array[1:] + self.q_array[:-1]) / 2

		self.qq_e_DiscreteN = qq_level/q_base

		return self.qq_e_DiscreteN

	def getSimuData(self):
		'''
		simulate key variables to generate targeted moments
		'''

		datalist = [[ self.y, self.n_opt[0], self.q_opt, self.q_array[1], self.qq_e_DiscreteN[1], self.qq_e_DiscreteN[2] ]]

		col_label = ["y", "n_opt", "q_opt", "q_n2", "qq_e_DN_2to3", "qq_e_DN_3to4"]

		return pd.DataFrame(data = datalist, columns = col_label)



class FamilyRation_ump():
	'''
	a simple version for fast simulation
	ump only, no decomposition
	choose among discrete n
	'''

	def __init__(self,
		pi_nq = 0.2,
		pi_n = 0.6, pi_q = 0, pi_s = 1,
		alpha = 0.35, rho = -6, theta = 0.5, gamma = 1,
		y = 10):

		# set up parameter values
		self.a = np.array([alpha, rho, theta, gamma])
		self.pi = np.array([pi_n, pi_q, pi_s, pi_nq])

		self.y = y
		# no target utility (tu in the Family() class)

		# initial values of the full range of commodities
		self.x = [1,1] # (q,s)
		self.n = [1]

		# setup model
		self.updateConstraints0()
		#self.setupModel()

	def evalUtility0(self, x, sign = 1.0):

		# theta: self.a[2]
		# alpha: self.a[0]
		# rho: self.a[1]
		# gamma: self.a[3]

		# pi_n: self.pi[0]
		# pi_q: self.pi[1]
		# pi_s: self.pi[2]
		# pi_nq: self.pi[3]

		# n: self.n[0]
		q = x[0]
		s = x[1]

		return sign * s**(1 - self.a[2]) * ( (self.a[0]*self.n[0]**self.a[1] + q**(self.a[3]*self.a[1])*(1 - self.a[0]))**(1/self.a[1]))**self.a[2]

	def evalUjac0(self, x, sign = 1.0):

		q = x[0]
		s = x[1]

		dudq = self.a[3]*q**(self.a[3]*self.a[1] - 1)*s**(1 - self.a[2])*self.a[2]*(1 - self.a[0])*((self.a[0]*self.n[0]**self.a[1] + 
			q**(self.a[3]*self.a[1])*(1 - self.a[0]))**(1/self.a[1]))**self.a[2]/((self.a[0]*self.n[0]**self.a[1] + q**(self.a[3]*self.a[1])*(1 - self.a[0])))

		duds = s**(- self.a[2])*(1 - self.a[2])*((self.a[0]*self.n[0]**self.a[1] + q**(self.a[3]*self.a[1])*(1 - self.a[0]))**(1/self.a[1]))**self.a[2] 
		return sign * np.array([dudq, duds])

	def evalExpenditure0(self, x):

		# theta: self.a[2]
		# alpha: self.a[0]
		# rho: self.a[1]
		# gamma: self.a[3]

		# pi_n: self.pi[0]
		# pi_q: self.pi[1]
		# pi_s: self.pi[2]
		# pi_nq: self.pi[3]

		# n: self.n[0]

		q = x[0]
		s = x[1]

		return self.n[0]*self.pi[0] + self.n[0]*self.pi[3]*q + self.pi[1]*q + self.pi[2]*s

	def evalEjac_m0(self, x):

		q = x[0]
		s = x[1]

		return -1.0 * np.array( [self.n[0] * self.pi[3] + self.pi[1], self.pi[2] ] )

	def excessWealth0(self, x):
		return self.y - self.evalExpenditure0(x)

	def updateConstraints0(self):
		'''
		Update the constraints for the UMP 
		'''
		self.cons_ump0 = (
			{
			'type':'eq',
			'fun': self.excessWealth0,
			'jac': self.evalEjac_m0
			},
			{
			'type':'ineq',
			'fun': lambda x: np.array(x),
			'jac': lambda x: np.identity(len(x))
			}
			)
	
	# ump, fast numeric version
	def ump_simple0(self, disp = False):
		'''
		The Utility Maximization Problem, simplified precedure
		'''
		if disp:
			print('-'*20)
			print('UMP Setup')
			print('Wealth: ' + str(self.y))
			print('Initial Commodity bunlde: ' + str(self.x))

		res = minimize(
			self.evalUtility0, self.x, args = (-1,),
			jac = self.evalUjac0,
			constraints = self.cons_ump0,
			method = 'SLSQP',
			options = {'disp':disp}
			)
		if res.success: 
			self.x = list(res.x)
			self.u = -res.fun
			if disp:
				print('UMP Results')
				print('Optimal Choice: {}'.format(res.x))
				print('Optimal Utility Level: {0}'.format(-res.fun) )
				print('-'*20)
		else: 
			print('UMP Failed! u = -9, x = [0,0]')
			print('res.x:', res.x)
			print('res.success:', res.success)			
			print('-'*20)
			self.x = [0,0]
			self.u = -9


	def evalDiscreteN0(self, disp = False):
		self.n_array = np.arange(1,7)
		len_n = len(self.n_array)
		self.u_array = np.zeros(len_n)
		self.q_array = np.zeros(len_n)
		
		if disp: print("Compute optimal fertility under discrete choices, n in 1 to 6 ")

		for i in range(0,len(self.n_array) ):
			
			if disp: print("Evaluate n = " + str(self.n_array[i]) )

			# set initial values of q0 and s0
			#q0 = self.a[0] * self.a[2] * net_y / ( self.pi[3] * self.n[0] )
			#s0 = (1 - self.a[2]) * net_y / self.pi[2] 
			#self.x = [q0/2, s0] 
			
			#self.x = [0.1 * net_y, 0.1 * net_y] 
			self.x = [0.1, 0.1] # an easy and fast initial value

			self.n = [ self.n_array[i] ]

			# pi_n: pi[0]; pi_nq: pi[3]; pi_s: pi[2]
			# theta: a[2]
			# alpha: a[0]

			# income net of expenses on n_bar
			net_y = self.y - self.pi[0] * self.n[0]
			if disp: print("Net y excluding expenses on n: " + str(net_y))
			if net_y < 0.10 * self.y:
				if disp: print("Net y < 0.10*y, pass")
				self.u_array[i] = -9 # negative u to denote negative net y
				continue

			# run ump
			#self.ump_simple(disp = disp)
			self.ump_simple0(disp = disp)

			# store utility and quality
			self.u_array[i] = self.u
			self.q_array[i] = self.x[0]
		
		i_max = np.argmax(self.u_array)

		self.n_opt = [ self.n_array[i_max] ]
		self.q_opt = self.q_array[i_max]

		return self.n_opt

	def evalQQ_DiscreteN(self):
		'''
		evaluate QQ effects in semi-elasticity form 
		after the execution of evalDiscreteN()
		based on the stored n_array and q_array
		'''

		# qq effect in level terms, n margins: 2->3, 3->4, 4->5
		qq_level = self.q_array[1:] - self.q_array[:-1]

		q_base = (self.q_array[1:] + self.q_array[:-1]) / 2

		self.qq_e_DiscreteN = qq_level/q_base

		return self.qq_e_DiscreteN

	def getSimuData(self):
		'''
		simulate key variables to generate targeted moments
		'''

		datalist = [[ self.y, self.n_opt[0], self.q_opt, self.q_array[1], self.qq_e_DiscreteN[1], self.qq_e_DiscreteN[2] ]]

		col_label = ["y", "n_opt", "q_opt", "q_n2", "qq_e_DN_2to3", "qq_e_DN_3to4"]

		return pd.DataFrame(data = datalist, columns = col_label)





